import os
import sys
import torch
import numpy as np
import random

# 添加项目根目录到路径
sys.path.append(os.path.abspath(os.path.dirname(__file__)))

from training.train_classifier_with_logging import train

# 设置随机种子以便结果可重现
def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

if __name__ == "__main__":
    # 设置随机种子
    set_seed(42)
    
    # 设置参数
    class Args:
        def __init__(self):
            # 数据目录 - 使用mini_face和mini_flower目录
            self.face_train = 'data/mini_face/train'  # 小型人脸训练数据
            self.face_val = 'data/mini_face/val'      # 小型人脸验证数据
            self.flower_train = 'data/mini_flower/train'  # 小型花卉训练数据
            self.flower_val = 'data/mini_flower/val'      # 小型花卉验证数据
            
            # 训练参数
            self.batch_size = 32  # 增加batch size以加快训练速度
            self.epochs = 5       # 5轮训练
            self.lr = 0.001
            
            # 模型输出目录结构 - 类似flower_model
            self.model_dir = 'output/face_flower_classifier'  # 模型主目录
            self.checkpoint_dir = os.path.join(self.model_dir, 'checkpoints')  # 检查点目录
            self.sample_dir = os.path.join(self.model_dir, 'samples')          # 样本图像目录
            self.tensorboard_dir = os.path.join(self.model_dir, 'logs')        # TensorBoard日志目录
            
            # 训练日志目录 - 存放在logs/classifier下
            self.log_dir = 'logs/classifier'
            
            self.num_workers = 2  # 减少worker数量
            self.no_cuda = False
    
    args = Args()
    
    # 创建输出目录结构
    os.makedirs(args.model_dir, exist_ok=True)
    os.makedirs(args.checkpoint_dir, exist_ok=True)
    os.makedirs(args.sample_dir, exist_ok=True)
    os.makedirs(args.tensorboard_dir, exist_ok=True)
    os.makedirs(args.log_dir, exist_ok=True)
    
    # 训练分类器
    print("开始训练人脸-花卉分类器...")
    print(f"使用数据集:")
    print(f"- 人脸训练: {args.face_train}")
    print(f"- 人脸验证: {args.face_val}")
    print(f"- 花卉训练: {args.flower_train}")
    print(f"- 花卉验证: {args.flower_val}")
    print(f"模型输出目录: {args.model_dir}")
    print(f"训练日志目录: {args.log_dir}")
    
    train(args)
    print("训练完成!")
    print(f"模型检查点已保存至: {args.checkpoint_dir}")
    print(f"训练日志已保存至: {args.log_dir}")
    print(f"TensorBoard日志: {args.tensorboard_dir}")
    print(f"样本图像: {args.sample_dir}") 